"""
钩子管理器模块，用于统一管理所有类型的钩子。
"""
from typing import Dict, List, Optional, Any, Callable, Union
import torch
from .gradient_hook import GradientHook
from .activation_hook import ActivationHook
from .optimizer_hook import OptimizerHook

class HookManager:
    """钩子管理器"""
    
    def __init__(self):
        """初始化钩子管理器"""
        self.gradient_hooks: Dict[str, GradientHook] = {}
        self.activation_hooks: Dict[str, ActivationHook] = {}
        self.optimizer_hooks: Dict[str, OptimizerHook] = {}
    
    def register_gradient_hook(self, 
                             name: str,
                             module: torch.nn.Parameter,
                             callback: Callable[[str, torch.Tensor, bool], None]) -> None:
        """
        注册梯度钩子
        
        Args:
            name: 参数名称
            module: 参数模块
            callback: 回调函数
        """
        if name in self.gradient_hooks:
            return
        
        hook = GradientHook(name, module, callback)
        self.gradient_hooks[name] = hook
    
    def register_activation_hook(self,
                               name: str,
                               module: torch.nn.Module,
                               callback: Callable[[str, torch.Tensor, bool], None]) -> None:
        """
        注册激活值钩子
        
        Args:
            name: 层名称
            module: 层模块
            callback: 回调函数
        """
        if name in self.activation_hooks:
            return
        
        hook = ActivationHook(name, module, callback)
        self.activation_hooks[name] = hook
    
    def register_optimizer_hook(self,
                              name: str,
                              optimizer: Any,
                              callback: Callable[[str, Dict[str, Any], float], None]) -> None:
        """
        注册优化器钩子
        
        Args:
            name: 优化器名称
            optimizer: 优化器实例
            callback: 回调函数
        """
        if name in self.optimizer_hooks:
            return
        
        hook = OptimizerHook(name, optimizer, callback)
        self.optimizer_hooks[name] = hook
    
    def remove_gradient_hook(self, name: str) -> None:
        """
        移除梯度钩子
        
        Args:
            name: 参数名称
        """
        if name in self.gradient_hooks:
            self.gradient_hooks[name].remove()
            del self.gradient_hooks[name]
    
    def remove_activation_hook(self, name: str) -> None:
        """
        移除激活值钩子
        
        Args:
            name: 层名称
        """
        if name in self.activation_hooks:
            self.activation_hooks[name].remove()
            del self.activation_hooks[name]
    
    def remove_optimizer_hook(self, name: str) -> None:
        """
        移除优化器钩子
        
        Args:
            name: 优化器名称
        """
        if name in self.optimizer_hooks:
            self.optimizer_hooks[name].remove()
            del self.optimizer_hooks[name]
    
    def remove_all_hooks(self) -> None:
        """移除所有钩子"""
        # 移除梯度钩子
        for name in list(self.gradient_hooks.keys()):
            self.remove_gradient_hook(name)
        
        # 移除激活值钩子
        for name in list(self.activation_hooks.keys()):
            self.remove_activation_hook(name)
        
        # 移除优化器钩子
        for name in list(self.optimizer_hooks.keys()):
            self.remove_optimizer_hook(name)
    
    def get_gradient_hooks(self) -> Dict[str, GradientHook]:
        """
        获取所有梯度钩子
        
        Returns:
            梯度钩子字典
        """
        return self.gradient_hooks
    
    def get_activation_hooks(self) -> Dict[str, ActivationHook]:
        """
        获取所有激活值钩子
        
        Returns:
            激活值钩子字典
        """
        return self.activation_hooks
    
    def get_optimizer_hooks(self) -> Dict[str, OptimizerHook]:
        """
        获取所有优化器钩子
        
        Returns:
            优化器钩子字典
        """
        return self.optimizer_hooks
    
    def __enter__(self):
        """上下文管理器入口"""
        return self
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        """上下文管理器出口"""
        self.remove_all_hooks() 